# 给你一个整数数组 nums 和一个整数 k ，请你统计并返回 该数组中和为 k 的子数组的个数 。
#
#  子数组是数组中元素的连续非空序列。
#
#  示例 1：
# 输入：nums = [1,1,1], k = 2
# 输出：2
#
#  示例 2：
# 输入：nums = [1,2,3], k = 3
# 输出：2
#
#
#  Related Topics 数组 哈希表 前缀和 👍 2329 👎 0

from typing import List, Dict
from collections import defaultdict


class Solution:
    def subarraySum2(self, nums: List[int], k: int) -> int:
        """
        前缀和 + dict
        """
        pre_sums_counts: Dict[int, int] = defaultdict(int)  # key前缀和，value为该和出现的次数
        res, pre_sum = 0, 0
        pre_sums_counts[0] = 1

        for num in nums:
            pre_sum = pre_sum + num
            res += pre_sums_counts[pre_sum - k]
            pre_sums_counts[pre_sum] += 1
        return res

    def subarraySum1(self, nums: List[int], k: int) -> int:
        """
        前缀和暴力法
        """
        pre_sums = [0] * (len(nums) + 1)
        res = 0
        for i in range(1, len(nums) + 1):
            pre_sums[i] = pre_sums[i - 1] + nums[i - 1]

        for i in range(1, len(nums) + 1):
            for j in range(i):
                if pre_sums[i] - pre_sums[j] == k:
                    res += 1
        return res

    def subarraySum(self, nums: List[int], k: int) -> int:
        return self.subarraySum2(nums, k)


# test code
if __name__ == "__main__":
    nums = [1, 1, 1]
    k = 2
    # nums = [1, 2, 3]
    # k = 3
    print(Solution().subarraySum(nums, k))
